"""
Feature Pyramid Network (FPN) Implementation

多尺度特征金字塔网络，用于处理不同尺度的目标。

参考：
- "Feature Pyramid Networks for Object Detection" (CVPR 2017)
- Detectron2实现
"""
import torch
import torch.nn as nn
import torch.nn.functional as F


class FPN(nn.Module):
    """
    Feature Pyramid Network

    输入：Backbone的多层特征 [C2, C3, C4, C5]
    输出：多尺度特征金字塔 [P2, P3, P4, P5]

    工作原理：
    1. 横向连接：1×1卷积统一通道数
    2. 自顶向下：从P5开始，逐层上采样并与横向特征相加
    3. 平滑：3×3卷积减少上采样的混叠效应
    """

    def __init__(self,
                 in_channels_list=[128, 384, 640, 640],  # OverLoCK各层通道数
                 out_channels=256,                       # FPN输出通道数
                 use_p2=True):                           # 是否使用P2层
        """
        Args:
            in_channels_list: 输入特征的通道数列表 [C2, C3, C4, C5]
            out_channels: FPN输出特征的统一通道数
            use_p2: 是否生成P2层（高分辨率但计算量大）
        """
        super(FPN, self).__init__()

        self.use_p2 = use_p2
        self.out_channels = out_channels

        # 横向连接：1×1卷积统一通道数
        self.lateral_convs = nn.ModuleList()
        for in_channels in in_channels_list:
            self.lateral_convs.append(
                nn.Conv2d(in_channels, out_channels, kernel_size=1)
            )

        # 平滑卷积：减少上采样的混叠效应
        self.fpn_convs = nn.ModuleList()
        num_fpn_layers = len(in_channels_list) if use_p2 else len(in_channels_list) - 1
        for _ in range(num_fpn_layers):
            self.fpn_convs.append(
                nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1)
            )

    def forward(self, features):
        """
        Args:
            features: 列表，包含多层特征 [C2, C3, C4, C5]
                     每个特征: [B, C_i, H_i, W_i]

        Returns:
            fpn_features: 列表，FPN特征金字塔 [P2, P3, P4, P5] 或 [P3, P4, P5]
                         每个特征: [B, out_channels, H_i, W_i]
        """
        # 检查输入
        assert len(features) == len(self.lateral_convs), \
            f"Expected {len(self.lateral_convs)} feature maps, got {len(features)}"

        # ===== 横向连接 =====
        # 将所有特征映射到统一通道数
        laterals = []
        for i, lateral_conv in enumerate(self.lateral_convs):
            laterals.append(lateral_conv(features[i]))

        # ===== 自顶向下路径 =====
        # 从最高层(C5)开始，逐层向下融合
        fpn_features = []

        # 最顶层：直接使用
        prev_features = laterals[-1]
        fpn_features.append(self.fpn_convs[-1](prev_features))

        # 从倒数第二层开始，向下融合
        for i in range(len(laterals) - 2, -1, -1):
            # 上采样前一层特征
            upsampled = F.interpolate(
                prev_features,
                size=laterals[i].shape[2:],
                mode='nearest'  # 最近邻插值（快速且效果好）
            )

            # 与横向特征相加
            prev_features = laterals[i] + upsampled

            # 平滑卷积
            fpn_features.append(self.fpn_convs[i](prev_features))

        # 反转列表（现在是自底向上：P2, P3, P4, P5）
        fpn_features = fpn_features[::-1]

        # 如果不使用P2，去掉最底层
        if not self.use_p2:
            fpn_features = fpn_features[1:]

        return fpn_features


class RoIFeatureExtractor(nn.Module):
    """
    基于FPN的RoI特征提取器

    根据ROI大小自动选择合适的特征层级，然后使用RoIAlign提取固定大小特征。
    """

    def __init__(self,
                 output_size=7,
                 sampling_ratio=2,
                 canonical_scale=224,
                 canonical_level=4,
                 fpn_levels=[2, 3, 4, 5]):
        """
        Args:
            output_size: RoIAlign输出大小 (output_size, output_size)
            sampling_ratio: RoIAlign采样点数（-1=自适应）
            canonical_scale: 标准尺度，用于分配ROI到不同层级
            canonical_level: 标准尺度对应的FPN层级
            fpn_levels: 可用的FPN层级
        """
        super(RoIFeatureExtractor, self).__init__()

        self.output_size = output_size
        self.canonical_scale = canonical_scale
        self.canonical_level = canonical_level
        self.fpn_levels = fpn_levels

        # 为每个FPN层级创建RoIAlign
        from torchvision.ops import RoIAlign

        self.roi_align_layers = nn.ModuleDict()
        for level in fpn_levels:
            stride = 2 ** level  # P2: stride=4, P3: stride=8, etc.
            self.roi_align_layers[f'p{level}'] = RoIAlign(
                output_size=(output_size, output_size),
                spatial_scale=1.0 / stride,
                sampling_ratio=sampling_ratio,
                aligned=True  # 重要！使用aligned版本
            )

    def _assign_boxes_to_levels(self, boxes, image_size):
        """
        根据ROI大小分配到不同的FPN层级

        公式：level = floor(canonical_level + log2(√(w*h) / canonical_scale))

        Args:
            boxes: [N, 4] ROI坐标 [x1, y1, x2, y2]
            image_size: (H, W) 原图尺寸

        Returns:
            levels: [N] 每个ROI对应的FPN层级
        """
        # 计算ROI面积
        widths = boxes[:, 2] - boxes[:, 0]
        heights = boxes[:, 3] - boxes[:, 1]
        areas = widths * heights

        # 计算目标层级
        # target_level = canonical_level + log2(sqrt(area) / canonical_scale)
        target_levels = self.canonical_level + torch.log2(torch.sqrt(areas) / self.canonical_scale)

        # 限制在有效范围内
        target_levels = torch.clamp(
            target_levels,
            min=min(self.fpn_levels),
            max=max(self.fpn_levels)
        )

        # 四舍五入到整数
        levels = torch.round(target_levels).to(torch.int64)

        return levels

    def forward(self, fpn_features, roi_boxes, image_size):
        """
        从FPN特征中提取ROI特征

        Args:
            fpn_features: 字典或列表，FPN特征金字塔
                         如果是列表: [P2, P3, P4, P5] 或 [P3, P4, P5]
                         如果是字典: {'p2': P2, 'p3': P3, 'p4': P4, 'p5': P5}
            roi_boxes: [N, 5] ROI坐标，格式 [batch_idx, x1, y1, x2, y2]
            image_size: (H, W) 原图尺寸

        Returns:
            roi_features: [N, C, output_size, output_size] ROI特征
        """
        # 将列表转换为字典
        if isinstance(fpn_features, list):
            fpn_dict = {}
            start_level = 2 if len(fpn_features) == 4 else 3
            for i, feat in enumerate(fpn_features):
                fpn_dict[f'p{start_level + i}'] = feat
            fpn_features = fpn_dict

        # 提取ROI坐标（去掉batch_idx）
        boxes_only = roi_boxes[:, 1:]  # [N, 4]

        # 分配ROI到不同层级
        levels = self._assign_boxes_to_levels(boxes_only, image_size)

        # 为每个层级提取特征
        roi_features_list = []
        roi_indices = []

        for level in self.fpn_levels:
            # 找到属于当前层级的ROI
            level_mask = levels == level

            if level_mask.sum() == 0:
                continue  # 跳过空层级

            # 提取当前层级的ROI
            level_boxes = roi_boxes[level_mask]  # [N_level, 5]

            # 获取对应的特征图
            level_features = fpn_features[f'p{level}']  # [B, C, H, W]

            # 使用RoIAlign提取
            roi_align = self.roi_align_layers[f'p{level}']
            level_roi_features = roi_align(level_features, level_boxes)  # [N_level, C, S, S]

            roi_features_list.append(level_roi_features)
            roi_indices.append(torch.where(level_mask)[0])

        # 合并所有层级的ROI特征，并恢复原始顺序
        roi_features = torch.cat(roi_features_list, dim=0)
        indices = torch.cat(roi_indices, dim=0)

        # 恢复原始顺序
        sorted_indices = torch.argsort(indices)
        roi_features = roi_features[sorted_indices]

        return roi_features


class FPNWithBackbone(nn.Module):
    """
    完整的FPN + Backbone模块

    集成backbone（如OverLoCK）和FPN，提供统一接口。
    """

    def __init__(self,
                 backbone,
                 backbone_out_indices=[0, 1, 2, 3],  # backbone输出哪几层
                 fpn_in_channels=[128, 384, 640, 640],
                 fpn_out_channels=256,
                 use_p2=True):
        """
        Args:
            backbone: Backbone网络（如OverLoCK）
            backbone_out_indices: backbone输出的层索引
            fpn_in_channels: 对应层的通道数
            fpn_out_channels: FPN输出通道数
            use_p2: 是否使用P2层
        """
        super(FPNWithBackbone, self).__init__()

        self.backbone = backbone
        self.backbone_out_indices = backbone_out_indices

        # FPN
        self.fpn = FPN(
            in_channels_list=fpn_in_channels,
            out_channels=fpn_out_channels,
            use_p2=use_p2
        )

        # ROI特征提取器
        fpn_levels = [2, 3, 4, 5] if use_p2 else [3, 4, 5]
        self.roi_extractor = RoIFeatureExtractor(
            output_size=7,
            fpn_levels=fpn_levels
        )

    def forward(self, images):
        """
        Args:
            images: [B, 3, H, W]

        Returns:
            fpn_features: FPN特征金字塔
        """
        # Backbone提取多层特征
        backbone_features = self.backbone.forward_features(images)

        # 选择需要的层
        if isinstance(backbone_features, list):
            selected_features = [backbone_features[i] for i in self.backbone_out_indices]
        else:
            # 如果backbone只返回单个特征，需要适配
            raise ValueError("Backbone must return list of features")

        # FPN
        fpn_features = self.fpn(selected_features)

        return fpn_features

    def extract_roi_features(self, fpn_features, roi_boxes, image_size):
        """
        从FPN中提取ROI特征

        Args:
            fpn_features: FPN特征金字塔
            roi_boxes: [N, 5] [batch_idx, x1, y1, x2, y2]
            image_size: (H, W)

        Returns:
            roi_features: [N, C, 7, 7]
        """
        return self.roi_extractor(fpn_features, roi_boxes, image_size)


def test_fpn():
    """测试FPN模块"""
    print("Testing FPN...")

    # 模拟OverLoCK的多层输出
    B = 2
    features = [
        torch.randn(B, 128, 56, 56),   # C2, stride=4
        torch.randn(B, 384, 28, 28),   # C3, stride=8
        torch.randn(B, 640, 14, 14),   # C4, stride=16
        torch.randn(B, 640, 7, 7),     # C5, stride=32
    ]

    # 创建FPN
    fpn = FPN(
        in_channels_list=[128, 384, 640, 640],
        out_channels=256,
        use_p2=True
    )

    # 前向传播
    fpn_features = fpn(features)

    print(f"\nFPN输出:")
    for i, feat in enumerate(fpn_features, start=2):
        print(f"  P{i}: {feat.shape}")

    # 测试ROI特征提取
    print("\nTesting RoI Feature Extractor...")

    roi_extractor = RoIFeatureExtractor(
        output_size=7,
        fpn_levels=[2, 3, 4, 5]
    )

    # 模拟8个ROI（4个来自图像0，4个来自图像1）
    roi_boxes = torch.tensor([
        [0, 10, 10, 50, 50],      # 小ROI → P2
        [0, 20, 20, 100, 100],    # 中ROI → P3
        [0, 50, 50, 200, 200],    # 大ROI → P4
        [0, 100, 100, 300, 300],  # 超大ROI → P5
        [1, 15, 15, 60, 60],
        [1, 30, 30, 120, 120],
        [1, 60, 60, 220, 220],
        [1, 110, 110, 310, 310],
    ], dtype=torch.float32)

    image_size = (224, 224)

    roi_features = roi_extractor(fpn_features, roi_boxes, image_size)

    print(f"ROI features shape: {roi_features.shape}")
    print(f"Expected: [8, 256, 7, 7]")

    assert roi_features.shape == (8, 256, 7, 7), "Shape mismatch!"

    print("\n✅ FPN test passed!")


if __name__ == '__main__':
    test_fpn()
